#include <Blas/spblas.cuh>
#include <cstring>

namespace myamg {

    __global__ void spmv8cols(const double* data, const int* colIndex, const double* x, double* y) {
        const unsigned int warpIdx = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
        const unsigned int laneIdx = threadIdx.x & 31;
        double fragA, fragB, acc[2];
        acc[0] = 0.0;
        acc[1] = 0.0;
        unsigned int idx = (warpIdx << 6) + laneIdx;
        fragA = loadGlobalWithCacheStreaming(data + idx);
        fragB = loadGlobalWithCacheStreaming(x + colIndex[idx]);
        FP64m8n8k4MMA(acc, fragA, fragB);
        idx += 32;
        fragA = loadGlobalWithCacheStreaming(data + idx);
        fragB = loadGlobalWithCacheStreaming(x + colIndex[idx]);
        FP64m8n8k4MMA(acc, fragA, fragB);
        const unsigned int mod9 = laneIdx % 9;
        if (mod9 == 0 || mod9 == 4) {
            const unsigned int outIdx = (laneIdx / 9) * 2 + (mod9 >> 2);
            storeGlobalWithCacheStreaming(y + MMA_M * warpIdx + outIdx, acc[mod9 >> 2]);
        }
    }

    void tensorSpMV(const int* Ap, const int* Ai, const double* Av, const double* x, double* y, const unsigned int nRows) {
        // 拷贝x的数据
        double* xDevice;
        cudaMalloc(&xDevice, nRows * sizeof(double));
        cudaMemcpy(xDevice, x, nRows * sizeof(double), cudaMemcpyHostToDevice);
        // 统计各种行的数量
        // short: 1-8个非零元
        // long: 9~12个非零元
        // very long: >12个非零元
        unsigned int numShortRows = 0;
        unsigned int numLongRows = 0;
        unsigned int numVeryLongRows = 0;
        for (unsigned int i = 0; i < nRows; ++i) {
            const auto numNoneZero = Ap[i + 1] - Ap[i];
            if (numNoneZero > 0) {
                if (numNoneZero <= 8) {
                    ++numShortRows;
                } else if (numNoneZero <= 12) {
                    ++numLongRows;
                } else {
                    ++numVeryLongRows;
                }
            }
        }
        // 建立映射
        auto* shortRowMap = new unsigned int[numShortRows];
        auto* longRowMap = new unsigned int[numLongRows];
        auto* veryLongRowMap = new unsigned int[numVeryLongRows];
        unsigned int shortIdx = 0;
        unsigned int longIdx = 0;
        unsigned int veryLongIdx = 0;
        for (unsigned int i = 0; i < nRows; ++i) {
            const auto numNoneZero = Ap[i + 1] - Ap[i];
            if (numNoneZero > 0) {
                if (numNoneZero <= 8) {
                    shortRowMap[shortIdx++] = i;
                } else if (numNoneZero <= 12) {
                    longRowMap[longIdx++] = i;
                } else {
                    veryLongRowMap[veryLongIdx++] = i;
                }
            }
        }
        // 处理short, 全部划分为2×8×4的块，不足者补零
        const auto shortRowPadded = alignUp<unsigned int>(numShortRows, 8);
        auto* shortDataHost = new double[shortRowPadded * 8 * 8];
        auto* shortColIndexHost = new int[shortRowPadded * 8 * 8];
        double* shortDataDevice;
        int* shortColIndexDevice;
        double* result;
        cudaMalloc(&shortDataDevice, shortRowPadded * 8 * 8 * sizeof(double));
        cudaMalloc(&shortColIndexDevice, shortRowPadded * 8 * 8 * sizeof(int));
        cudaMallocManaged(&result, shortRowPadded * sizeof(double));
        // 填充零
        std::memset(shortDataHost, 0, shortRowPadded * 8 * 8 * sizeof(double));
        std::memset(shortColIndexHost, 0, shortRowPadded * 8 * 8 * sizeof(int));
        // 遍历所有short行,填充数据，每八行为一个单位进行处理
        // 内存布局如下
        // |行0的前四列|行1的前四列|...|行7的前四列|行0的后四列|行1的后四列|...|行7的后四列|
        for (auto i = 0; i < (shortRowPadded / 8); ++i) {
            // 处理第i个8行块，offset是当前行块的第一行
            const auto offset = i * 8;
            // 处理8行块的8行
            for (auto j = 0; j < 8; ++j) {
                // 实际行号
                const auto actualRow = j + offset;
                if (actualRow >= numShortRows) {
                    break;
                }
                // 获取真实行号
                auto realRow = shortRowMap[actualRow];
                // 处理该行的非零元
                auto numNoneZero = Ap[realRow + 1] - Ap[realRow];
                for (auto k = 0; k < 8; ++k) {
                    if (k >= numNoneZero) {
                        break;
                    }
                    // 填充数据， 前四列为一组，后四列为一组
                    const auto idx = i * 64 + (k >> 2) * 32 + j * 4 + (k & 3);
                    shortDataHost[idx] = Av[Ap[realRow] + k];
                    shortColIndexHost[idx] = Ai[Ap[realRow] + k];
                }
            }
        }
        // 拷贝到设备
        cudaMemcpy(shortDataDevice, shortDataHost, shortRowPadded * 8 * 8 * sizeof(double), cudaMemcpyHostToDevice);
        cudaMemcpy(shortColIndexDevice, shortColIndexHost, shortRowPadded * 8 * 8 * sizeof(int), cudaMemcpyHostToDevice);
        // 调用核函数
        dim3 block(512);
        const auto numWarps = shortRowPadded / 8;
        const auto WarpPerBlock = block.x / 32;
        dim3 grid((numWarps + WarpPerBlock - 1) / WarpPerBlock);
        spmv8cols<<<grid, block>>>(shortDataDevice, shortColIndexDevice, xDevice, result);
        cudaDeviceSynchronize();
        // 将结果写回y
        for (auto i = 0; i < numShortRows; ++i) {
            const auto realRow = shortRowMap[i];
            y[realRow] = result[i];
        }
        // 测试结果正确性，long行和very long行直接在CPU上计算
        for (auto i = 0; i < numLongRows; ++i) {
            const auto realRow = longRowMap[i];
            double sum = 0.0;
            for (auto j = Ap[realRow]; j < Ap[realRow + 1]; ++j) {
                sum += Av[j] * x[Ai[j]];
            }
            y[realRow] = sum;
        }
        for (auto i = 0; i < numVeryLongRows; ++i) {
            const auto realRow = veryLongRowMap[i];
            double sum = 0.0;
            for (auto j = Ap[realRow]; j < Ap[realRow + 1]; ++j) {
                sum += Av[j] * x[Ai[j]];
            }
            y[realRow] = sum;
        }
        // 释放资源
        cudaFree(xDevice);
        cudaFree(shortDataDevice);
        cudaFree(shortColIndexDevice);
        cudaFree(result);
        delete[] shortRowMap;
        delete[] longRowMap;
        delete[] veryLongRowMap;
        delete[] shortDataHost;
        delete[] shortColIndexHost;
    }

}